import serial
import minimalmodbus
import time
import threading
import sys
import struct
from datetime import datetime  # 导入时间模块

class HaitaiMotor:
    def __init__(self, serial_port, motor_id=0x01, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.ser = serial_port
        self.motor_id = motor_id  # 协议默认设备地址0x01
        self.lock = lock or threading.Lock()
        self.counts_per_circle = 16384  # 协议定义：一圈=16384 Counts
        self.packet_seq = 0x00  # 包序号，初始为0x00

    def _crc16_modbus(self, data):
        """计算CRC16_MODBUS校验码（与C++代码逻辑完全一致）"""
        crc = 0xFFFF
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc & 0xFFFF

    def _pack_frame(self, cmd, data_bytes=b''):
        """按协议格式打包帧（对齐C++代码的命令结构）"""
        frame_header = [
            0x3E,  # 协议头（主机发送）
            self.packet_seq,  # 包序号
            self.motor_id,  # 设备地址
            cmd,  # 命令码
            len(data_bytes)  # 数据长度
        ]
        frame_header_bytes = bytes(frame_header)
        # 计算CRC16（协议头至数据字段）
        crc = self._crc16_modbus(frame_header_bytes + data_bytes)
        crc_bytes = struct.pack('<H', crc)  # 小端字节序（与C++一致）
        # 完整帧
        full_frame = frame_header_bytes + data_bytes + crc_bytes
        # 包序号自增（0x00-0xFF循环）
        self.packet_seq = (self.packet_seq + 1) & 0xFF
        return full_frame

    def _send_command(self, cmd, data_bytes=b'', read_len=15, post_delay=0.1):
        """发送命令并接收应答（参考C++的延迟和缓冲区处理）"""
        full_frame = self._pack_frame(cmd, data_bytes)
        with self.lock:
            try:
                # 清空缓冲区（与C++ clearSerialBuffer功能一致）
                self.ser.reset_input_buffer()
                self.ser.reset_output_buffer()
                self.ser.write(full_frame)
                self.ser.flush()
                time.sleep(post_delay)  # 延长延迟，确保电机应答完成
                resp = self.ser.read(read_len)
                # 验证应答协议头（应返回0x3C）
                if resp and len(resp) >= 5 and resp[0] == 0x3C:
                    return resp
                else:
                    print(f"[电机] 应答无效: 长度={len(resp) if resp else 0}, 数据={resp.hex() if resp else '无'}")
                    return None
            except Exception as e:
                print(f"[电机] 发送命令失败: {e}")
                raise

    def get_current_angle(self):
        """读取当前角度（完全对齐C++的getMotorAngleValues逻辑）"""
        try:
            # 命令码0x2F，读取角度（C++同款命令）
            resp = self._send_command(0x2F, b'', read_len=15, post_delay=0.1)
        except Exception as e:
            print(f"[电机] 读取角度出错: {e}")
            return None
        
        if resp and len(resp) == 15:  # C++要求应答长度必须为15字节
            try:
                # 1. 解析单圈绝对值角度（uint16_t，与C++一致）
                single_turn_raw = (resp[5] | (resp[6] << 8)) & 0x3FFF  # 掩码0x3FFF（协议要求）
                single_turn_angle = single_turn_raw * (360.0 / self.counts_per_circle)
                
                # 2. 解析多圈绝对值角度（int32_t，与C++一致的字节拼接逻辑）
                multi_turn_raw = (resp[7] | (resp[8] << 8) | (resp[9] << 16) | (resp[10] << 24))
                # 转换为Python的int32类型（处理符号位）
                if multi_turn_raw & 0x80000000:
                    multi_turn_raw -= 0x100000000
                multi_turn_angle = multi_turn_raw * (360.0 / self.counts_per_circle)
                
                # 返回多圈总角度（与C++逻辑一致）
                return multi_turn_angle
            except Exception as e:
                print(f"[电机] 解析角度失败: {e} | 应答数据: {resp.hex()}")
                return None
        else:
            print(f"[电机] 应答长度错误（需15字节）: 实际={len(resp) if resp else 0}")
            return None

    def set_origin(self):
        """设置当前位置为原点（命令码0x21）"""
        print("设置当前位置为原点...")
        try:
            resp = self._send_command(0x21, b'', read_len=10, post_delay=0.5)
            if resp and len(resp) >= 10:
                result = resp[7]  # 0x01成功，0x00失败（协议要求）
                if result == 0x01:
                    print("原点设置成功")
                    return True
                else:
                    print(f"原点设置失败，结果码: 0x{result:02X}")
                    return False
            else:
                print(f"[电机] 设置原点应答无效: {resp.hex() if resp else '无响应'}")
                return False
        except Exception as e:
            print(f"[电机] 设置原点失败: {e}")
            return False

    def move_to_absolute_position(self, target_angle, max_speed=1000):
        """绝对值位置闭环控制（命令码0x55）"""
        # 角度转Counts（与C++一致的计算逻辑）
        target_counts = int(target_angle * (self.counts_per_circle / 360.0))
        # 目标位置为uint32_t（4字节小端，与协议一致）
        target_bytes = struct.pack('<I', target_counts)
        # 配置位置闭环目标速度（命令码0x57）
        speed_bytes = struct.pack('<H', max_speed)  # 速度单位0.1RPM，int16_t
        try:
            # 先配置目标速度（0x01=配置模式）
            self._send_command(0x57, b'\x01' + speed_bytes, post_delay=0.2)
            time.sleep(0.1)
            # 发送绝对值位置控制命令
            self._send_command(0x55, target_bytes, post_delay=0.2)
            return True
        except Exception as e:
            print(f"[电机] 绝对值位置控制失败: {e}")
            return False

    def move_to_zero(self, timeout=300):
        """回归零位（设置原点后定位到0度）"""
        print("开始回归零位...")
        # 1. 设置当前位置为原点
        if not self.set_origin():
            return False
        # 2. 定位到0度（绝对值位置控制）
        if not self.move_to_absolute_position(0, max_speed=500):  # 500=50.0 0.1RPM
            return False
        # 3. 等待到达目标位置
        start_time = time.time()
        while time.time() - start_time < timeout:
            current_angle = self.get_current_angle()
            if current_angle is not None and abs(current_angle) <= 0.1:
                print("已到达零位")
                return True
            time.sleep(0.2)
        print("回零超时")
        return False

    def rotate_360(self, timeout=300):
        """旋转360度（绝对值位置控制）"""
        print("开始旋转360度...")
        # 定位到360度（对应Counts=16384）
        if not self.move_to_absolute_position(360, max_speed=20):  # 1000=100.0 0.1RPM
            return False
        # 等待旋转完成
        start_time = time.time()
        while time.time() - start_time < timeout:
            current_angle = self.get_current_angle()
            if current_angle is not None:
                normalized_angle = current_angle % 360
                if abs(normalized_angle - 0) <= 0.1:
                    print("已完成360度旋转")
                    return True
            time.sleep(0.2)
        print("旋转超时")
        return False

    def read_status(self):
        """读取电机状态（命令码0x40，修复温度解析错误）"""
        try:
            resp = self._send_command(0x40, b'', read_len=12, post_delay=0.1)
            if resp and len(resp) >= 12:
                # 修复：电压/电流/温度为1字节（协议V2.3明确规定，之前误读为2字节）
                voltage = resp[5] * 0.2  # 1字节，0.2V/单位（协议第10页）
                current = resp[6] * 0.03  # 1字节，0.03A/单位（协议第10页）
                temperature = resp[7] * 0.4  # 1字节，0.4℃/单位（协议第10页）
                fault_code = resp[8]
                run_mode = resp[9]
                mode_map = {0: "关闭状态", 1: "开环模式", 3: "速度模式", 5: "位置模式"}
                print(f"[状态] 电压: {voltage:.1f}V, 电流: {current:.2f}A, 温度: {temperature:.1f}℃, 故障码: 0x{fault_code:02X}, 运行模式: {mode_map.get(run_mode, '未知')}")
                return {
                    'voltage': voltage,
                    'current': current,
                    'temperature': temperature,
                    'fault_code': fault_code,
                    'run_mode': run_mode
                }
            else:
                print(f"[电机] 读取状态应答无效: {resp.hex() if resp else '无响应'}")
        except Exception as e:
            print(f"[电机] 读取状态失败: {e}")
        return None

    def close_motor(self):
        """关闭电机（命令码0x50，与C++ sendMotorShutdownCommand一致）"""
        try:
            self._send_command(0x50, b'', post_delay=0.2)
            print("电机已关闭，进入自由态")
        except Exception as e:
            print(f"[电机] 关闭失败: {e}")

    def close(self):
        self.close_motor()

class LE100485Sensor:
    def __init__(self, serial_port, sensor_address=2, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.serial_port = serial_port
        self.lock = lock or threading.Lock()
        self.sensor_address = sensor_address
        self.instrument = minimalmodbus.Instrument(serial_port, slaveaddress=self.sensor_address)
        self.instrument.mode = minimalmodbus.MODE_RTU
        self.instrument.serial = self.serial_port
        self.instrument.close_port_after_each_call = False
        self.instrument.clear_buffers_before_each_transaction = True
        self.instrument.debug = False
        self.distance_register = 0x003B
        self.resolution = 0.001

    def read_distance(self, register_address=None):
        if register_address is None:
            register_address = self.distance_register
        with self.lock:
            try:
                raw_data = self.instrument.read_register(
                    registeraddress=register_address,
                    functioncode=3,
                    signed=False
                )
                return raw_data * self.resolution
            except Exception as e:
                print(f"[Sensor] 读取失败: {e}")
                return None

    def close(self):
        pass

class DataCollector:
    def __init__(self, motor, sensor):
        self.motor = motor
        self.sensor = sensor
        self.data = []
        self.running = False
        self.thread = None

    def _collect_loop(self, interval):
        while self.running:
            angle = self.motor.get_current_angle()
            distance = self.sensor.read_distance()
            if angle is not None and distance is not None:
                normalized_angle = angle % 360
                # 新增：记录当前采集时间（精确到秒）
                collect_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                self.data.append((collect_time, normalized_angle, distance))  # 新增时间字段
                print(f"时间: {collect_time} | 角度: {normalized_angle:.2f}°  距离: {distance:.5f}mm", end='\r', flush=True)
            time.sleep(interval)

    def start(self, interval=0.05):
        if self.running:
            return
        self.data = []
        self.running = True
        self.thread = threading.Thread(target=self._collect_loop, args=(interval,), daemon=True)
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join(timeout=2)

    def save_data(self, filename=None):
        # 若未指定文件名，按当前时间（精确到秒）生成CSV文件名
        if filename is None:
            current_time = datetime.now().strftime("%Y%m%d_%H%M%S")  # 格式：年-月-日_时-分-秒
            filename = f"measurement_data_{current_time}.csv"  # 后缀改为csv
        
        # 保存为CSV格式（逗号分隔，兼容Excel）
        with open(filename, 'w', encoding='utf-8-sig') as f:  # utf-8-sig支持中文编码
            # 新增时间字段表头
            f.write("采集时间,角度(度),距离(mm)\n")
            for collect_time, angle, distance in self.data:
                f.write(f"{collect_time},{angle:.4f},{distance:.5f}\n")
        print(f"\n数据已保存到 {filename}，共 {len(self.data)} 条记录")

def main():
    # 配置参数
    PORT = 'COM6'  
    BAUDRATE = 115200
    MOTOR_ID = 0x01  # 默认设备地址
    SENSOR_ADDRESS = 2
    shared_serial = None
    collector = None

    try:
        shared_serial = serial.Serial(
            port=PORT,
            baudrate=BAUDRATE,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS,
            timeout=1,
            write_timeout=5
        )

        # 启用RS485模式（适配协议要求）
        try:
            if hasattr(serial, "rs485"):
                rs485_settings = serial.rs485.RS485Settings(
                    rts_level_for_tx=True,
                    rts_level_for_rx=False,
                    delay_before_tx=0.002,
                    delay_before_rx=0.002
                )
                shared_serial.rs485_mode = rs485_settings
                print("[Main] 已启用RS485模式")
        except Exception as e:
            print(f"[Main] 启用RS485模式失败（忽略）: {e}")

        serial_lock = threading.Lock()
        motor = HaitaiMotor(serial_port=shared_serial, motor_id=MOTOR_ID, lock=serial_lock)
        sensor = LE100485Sensor(serial_port=shared_serial, sensor_address=SENSOR_ADDRESS, lock=serial_lock)
        collector = DataCollector(motor, sensor)

        # 读取电机初始状态
        motor.read_status()

        # 回归零位
        if not motor.move_to_zero():
            print("无法回归零位，程序终止")
            return

        # 开始数据采集
        collector.start(interval=0.05)

        # 旋转360度
        if not motor.rotate_360():
            print("旋转失败，停止采集")
            collector.stop()
            return

        # 停止采集并保存数据（不指定文件名，自动按时间生成）
        collector.stop()
        collector.save_data()  # 无需传参，自动生成时间命名的CSV

        # 回归零位并关闭电机
        motor.move_to_zero()
        motor.close_motor()
        print("测量流程完成")

    except serial.SerialException as e:
        print(f"串口通信错误: {e} | 请检查：1.串口号是否为{PORT} 2.USB转485是否连接 3.电机是否供电")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        if collector:
            collector.stop()
        if shared_serial and shared_serial.is_open:
            shared_serial.close()
            print("串口已关闭")

if __name__ == "__main__":
    main()